import torch
import torch.nn as nn
import torch.nn.functional as F

"""
在 forward 函数定义中的 * 符号有以下作用：
强制关键字参数
* 之后的参数必须以关键字形式传递
不能通过位置参数的方式传入
"""
test_input = torch.randint(0, 10, (4, 8))

key_padding_mask = torch.Tensor([[True, True, True, True, True, True, True, True],
                                 [True, True, True, True, True, True, True, False],
                                 [True, True, True, True, True, True, False, False],
                                 [True, True, True, True, False, False, False, False]])


class MaskedMHA(nn.Module):
    def __init__(self, d_model=512, num_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0
        self.h = num_heads
        # 执行整数除法：
        # // 将两个数相除并向下取整到最接近的整数
        # 结果始终是整数类型（即使操作数是浮点数）
        self.d = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)

        self.W_o = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x):  # [B,L,D] -> [B,h,L,d]
        B, L, _ = x.shape
        return x.view(B, L, self.h, self.d).transpose(1, 2).contiguous()

    def forward(self, x, *, causal=True, key_padding_mask=None):
        """
        x: [B, L, D]  （自注意力：Q=K=V=x）
        key_padding_mask: [B, L]，True=有效/非pad，False=pad
        return y: [B, L, D], attn: [B, h, L, L]
        """
        B, L, D = x.shape
        Q = self._split_heads(self.W_q(x))
        K = self._split_heads(self.W_k(x))
        V = self._split_heads(self.W_v(x))

        # 1) 注意力分数
        scores = (Q @ K.transpose(-2, -1)) / (self.d ** 0.5)  # [B,h,L,L]
        print("scores:", scores)

        # 2) 因果掩码（下三角保留）
        if causal:
            # torch.tril()创建下三角矩阵作为因果掩码，确保每个位置只能关注到它之前的位置
            # masked_fill()将掩码为False的位置（即未来位置）的注意力分数设为负无穷，这样在softmax后这些位置的注意力权重就变为0
            causal_mask = torch.tril(torch.ones(L, L, device=x.device, dtype=torch.bool))
            print("causal_mask:", causal_mask)
            scores = scores.masked_fill(~causal_mask, float('-inf'))
            print("scores:", scores)

        # 3) key padding 掩码：把 K/V 的 pad 位置遮掉
        if key_padding_mask is not None:  # [B,L] → [B,1,1,L]
            """
            1. 数据类型不确定性
            虽然在您的测试代码中 key_padding_mask 被定义为 torch.Tensor 并包含布尔值
            但实际上 torch.Tensor 默认创建的是 float32 类型，而非布尔类型
            即使看起来像布尔值，它们在PyTorch内部仍可能是浮点数（如1.0和0.0）
            2. 确保类型一致性
            .bool() 显式地将张量转换为布尔类型
            这确保了在后续的 masked_fill 操作中，掩码条件判断的可靠性
            避免由于类型不匹配可能导致的意外行为
            """
            print("key_padding_mask:", key_padding_mask)
            kmask = key_padding_mask.unsqueeze(1).unsqueeze(1).bool()
            print("kmask:", kmask)
            scores = scores.masked_fill(~kmask, float('-inf'))
            print("scores:", scores)

        # 4) softmax + dropout
        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)

        # 5) 加权求和 + 合并头
        ctx = attn @ V  # [B,h,L,d]
        ctx = ctx.transpose(1, 2).reshape(B, L, D)  # [B,L,D]
        y = self.W_o(ctx)
        return y, attn


class MaskedMHA_test(nn.Module):
    def __init__(self):
        super(MaskedMHA_test, self).__init__()
        self.embedding = nn.Embedding(10, 512)
        self.masked_mha = MaskedMHA(512, 8, 0.1)
        self.linear = nn.Linear(512, 10)

    def forward(self, x):
        x = self.embedding(x)
        x, attn = self.masked_mha(x, causal=True, key_padding_mask=key_padding_mask)
        x = self.linear(x)
        return x


if __name__ == '__main__':
    model = MaskedMHA_test()
    print(model(test_input))
